"""
音频压缩和格式转换模块
提供高效的音频压缩和格式转换功能
"""

import io
import os
import tempfile
from typing import Dict, Any, Optional, Tuple, Union
from dataclasses import dataclass
from enum import Enum
import numpy as np
import soundfile as sf
from pydub import AudioSegment
import librosa


class CompressionQuality(Enum):
    """压缩质量等级"""
    LOW = "low"       # 64k
    MEDIUM = "medium"  # 128k
    HIGH = "high"     # 192k
    ULTRA = "ultra"   # 320k


@dataclass
class CompressionSettings:
    """压缩设置"""
    quality: CompressionQuality = CompressionQuality.MEDIUM
    format: str = 'mp3'
    sample_rate: Optional[int] = None
    channels: Optional[int] = None
    bitrate: Optional[str] = None
    normalize: bool = True
    trim_silence: bool = False
    apply_limiter: bool = True


@dataclass
class CompressionResult:
    """压缩结果"""
    success: bool
    compressed_data: Optional[bytes] = None
    original_size: int = 0
    compressed_size: int = 0
    compression_ratio: float = 0.0
    duration: Optional[float] = None
    quality_metrics: Optional[Dict[str, Any]] = None
    error_message: Optional[str] = None


class AudioCompressor:
    """音频压缩器"""
    
    def __init__(self):
        self.quality_bitrates = {
            CompressionQuality.LOW: "64k",
            CompressionQuality.MEDIUM: "128k",
            CompressionQuality.HIGH: "192k",
            CompressionQuality.ULTRA: "320k"
        }
        
        self.supported_formats = {
            'mp3', 'wav', 'flac', 'ogg', 'aac', 'm4a'
        }
    
    def compress_audio(
        self,
        input_data: Union[bytes, np.ndarray],
        sample_rate: int,
        settings: CompressionSettings
    ) -> CompressionResult:
        """压缩音频"""
        try:
            # 计算原始大小
            if isinstance(input_data, bytes):
                original_size = len(input_data)
            else:
                original_size = input_data.nbytes
            
            # 处理音频数据
            processed_data, processed_sr = self._preprocess_audio(
                input_data, sample_rate, settings
            )
            
            # 执行压缩
            compressed_data = self._apply_compression(
                processed_data, processed_sr, settings
            )
            
            # 计算压缩结果
            compressed_size = len(compressed_data)
            compression_ratio = compressed_size / original_size if original_size > 0 else 0
            
            # 计算质量指标
            duration = len(processed_data) / processed_sr
            quality_metrics = self._calculate_quality_metrics(
                input_data, sample_rate, compressed_data, settings
            )
            
            return CompressionResult(
                success=True,
                compressed_data=compressed_data,
                original_size=original_size,
                compressed_size=compressed_size,
                compression_ratio=compression_ratio,
                duration=duration,
                quality_metrics=quality_metrics
            )
            
        except Exception as e:
            return CompressionResult(
                success=False,
                error_message=f"Compression failed: {str(e)}"
            )
    
    def _preprocess_audio(
        self,
        input_data: Union[bytes, np.ndarray],
        sample_rate: int,
        settings: CompressionSettings
    ) -> Tuple[np.ndarray, int]:
        """预处理音频"""
        # 如果输入是字节数据，先转换为numpy数组
        if isinstance(input_data, bytes):
            data, sr = self._bytes_to_array(input_data)
        else:
            data = input_data
            sr = sample_rate
        
        # 重采样
        if settings.sample_rate and settings.sample_rate != sr:
            data = librosa.resample(data, orig_sr=sr, target_sr=settings.sample_rate)
            sr = settings.sample_rate
        
        # 声道转换
        if settings.channels == 1 and len(data.shape) > 1:
            data = np.mean(data, axis=1)
        elif settings.channels == 2 and len(data.shape) == 1:
            data = np.column_stack([data, data])
        
        # 标准化
        if settings.normalize:
            data = self._normalize_audio(data)
        
        # 去除静音
        if settings.trim_silence:
            data, _ = librosa.effects.trim(data, top_db=20)
        
        # 应用限幅器
        if settings.apply_limiter:
            data = self._apply_limiter(data)
        
        return data, sr
    
    def _apply_compression(
        self,
        data: np.ndarray,
        sample_rate: int,
        settings: CompressionSettings
    ) -> bytes:
        """应用压缩"""
        # 确定比特率
        if settings.bitrate:
            bitrate = settings.bitrate
        else:
            bitrate = self.quality_bitrates[settings.quality]
        
        # 使用pydub进行压缩
        audio = self._array_to_segment(data, sample_rate)
        
        # 创建内存缓冲区
        output_buffer = io.BytesIO()
        
        # 导出压缩音频
        export_params = {
            'format': settings.format,
            'bitrate': bitrate
        }
        
        if settings.format in ['mp3', 'aac', 'm4a']:
            # 有损格式
            audio.export(output_buffer, **export_params)
        else:
            # 无损格式
            export_params.pop('bitrate')
            audio.export(output_buffer, **export_params)
        
        output_buffer.seek(0)
        return output_buffer.read()
    
    def _bytes_to_array(self, audio_bytes: bytes) -> Tuple[np.ndarray, int]:
        """将字节数据转换为numpy数组"""
        with tempfile.NamedTemporaryFile() as temp_file:
            temp_file.write(audio_bytes)
            temp_file.seek(0)
            
            data, sr = sf.read(temp_file.name)
            return data, sr
    
    def _array_to_segment(self, data: np.ndarray, sample_rate: int) -> AudioSegment:
        """将numpy数组转换为AudioSegment"""
        # 转换为16位整数
        data_int16 = (data * np.iinfo(np.int16).max).astype(np.int16)
        
        # 创建AudioSegment
        audio = AudioSegment(
            data_int16.tobytes(),
            frame_rate=sample_rate,
            sample_width=2,  # 16-bit
            channels=1 if len(data.shape) == 1 else data.shape[1]
        )
        
        return audio
    
    def _normalize_audio(self, data: np.ndarray) -> np.ndarray:
        """标准化音频"""
        if len(data) == 0:
            return data
        
        # 计算RMS
        rms = np.sqrt(np.mean(data ** 2))
        
        if rms > 0:
            # 标准化到-12dB
            target_rms = 0.25
            data = data * (target_rms / rms)
        
        return np.clip(data, -1.0, 1.0)
    
    def _apply_limiter(self, data: np.ndarray, threshold: float = -0.5) -> np.ndarray:
        """应用限幅器"""
        # 简单的软限幅器
        threshold_linear = np.power(10, threshold / 20)
        
        # 应用软限幅
        limited_data = np.where(
            np.abs(data) > threshold_linear,
            threshold_linear * np.tanh(data / threshold_linear),
            data
        )
        
        return limited_data
    
    def _calculate_quality_metrics(
        self,
        original_data: Union[bytes, np.ndarray],
        sample_rate: int,
        compressed_data: bytes,
        settings: CompressionSettings
    ) -> Dict[str, Any]:
        """计算质量指标"""
        try:
            # 如果原始数据是字节数组，转换为numpy数组
            if isinstance(original_data, bytes):
                original_array, _ = self._bytes_to_array(original_data)
            else:
                original_array = original_data
            
            # 解压缩数据
            compressed_array, _ = self._bytes_to_array(compressed_data)
            
            # 确保两个数组长度相同
            min_length = min(len(original_array), len(compressed_array))
            if min_length == 0:
                return {"error": "Empty audio data"}
            
            original_array = original_array[:min_length]
            compressed_array = compressed_array[:min_length]
            
            # 计算SNR
            signal_power = np.mean(original_array ** 2)
            noise_power = np.mean((original_array - compressed_array) ** 2)
            snr = 10 * np.log10(signal_power / noise_power) if noise_power > 0 else np.inf
            
            # 计算MSE
            mse = np.mean((original_array - compressed_array) ** 2)
            
            # 计算PSNR
            max_signal = np.max(np.abs(original_array))
            psnr = 20 * np.log10(max_signal / np.sqrt(mse)) if mse > 0 else np.inf
            
            return {
                'snr_db': float(snr),
                'mse': float(mse),
                'psnr_db': float(psnr),
                'bitrate': settings.bitrate or self.quality_bitrates[settings.quality],
                'format': settings.format
            }
            
        except Exception as e:
            return {"error": f"Quality calculation failed: {str(e)}"}
    
    def batch_compress(
        self,
        audio_list: list,
        sample_rate: int,
        settings: CompressionSettings
    ) -> list[CompressionResult]:
        """批量压缩音频"""
        results = []
        
        for i, audio_data in enumerate(audio_list):
            result = self.compress_audio(audio_data, sample_rate, settings)
            result.metadata = {"index": i}
            results.append(result)
        
        return results
    
    def compress_file(
        self,
        input_path: str,
        output_path: str,
        settings: CompressionSettings
    ) -> CompressionResult:
        """压缩音频文件"""
        try:
            # 读取音频文件
            data, sr = sf.read(input_path)
            
            # 压缩音频
            result = self.compress_audio(data, sr, settings)
            
            # 保存压缩后的文件
            if result.success and result.compressed_data:
                with open(output_path, 'wb') as f:
                    f.write(result.compressed_data)
                
                result.output_path = output_path
            
            return result
            
        except Exception as e:
            return CompressionResult(
                success=False,
                error_message=f"File compression failed: {str(e)}"
            )
    
    def get_compression_estimate(
        self,
        duration: float,
        sample_rate: int,
        channels: int,
        settings: CompressionSettings
    ) -> Dict[str, Any]:
        """估算压缩结果"""
        # 估算原始大小
        original_size = int(duration * sample_rate * channels * 2)  # 16-bit audio
        
        # 估算压缩后大小
        if settings.bitrate:
            bitrate_kbps = int(settings.bitrate.replace('k', ''))
        else:
            bitrate_kbps = int(self.quality_bitrates[settings.quality].replace('k', ''))
        
        compressed_size = int(duration * bitrate_kbps * 1000 / 8)  # 字节
        
        # 估算压缩比
        compression_ratio = compressed_size / original_size if original_size > 0 else 0
        
        return {
            'original_size_mb': original_size / (1024 * 1024),
            'compressed_size_mb': compressed_size / (1024 * 1024),
            'compression_ratio': compression_ratio,
            'space_saved_mb': (original_size - compressed_size) / (1024 * 1024),
            'estimated_quality': settings.quality.value
        }
    
    def create_preset(self, name: str, settings: CompressionSettings):
        """创建压缩预设"""
        # 这里可以实现预设管理
        pass
    
    def get_supported_formats(self) -> list[str]:
        """获取支持的格式列表"""
        return list(self.supported_formats)
